#include <jni.h>
#include <string>
#include <vector>
#include <cassert>

#include <android/log.h>
#include <android/native_window_jni.h>
#include <vulkan/vulkan.h>
#include <vulkan/vulkan_android.h>
#include <vulkan/vulkan_core.h>

extern "C" JNIEXPORT jstring

JNICALL
Java_com_june_myfps1_MainActivity_stringFromJNI(
        JNIEnv *env,
        jobject /* this */) {
    std::string hello = "Hello from my fps1";
    return env->NewStringUTF(hello.c_str());
}

struct VulkanContext {
    VkInstance instance;
    VkSurfaceKHR surface;
    VkPhysicalDevice physicalDevice;
    VkDevice device;

    VkQueue graphicsQueue;
    VkQueue presentQueue;
    uint32_t graphicsQueueFamily;
    uint32_t presentQueueFamily;

    VkSwapchainKHR swapchain;
    std::vector<VkImage> swapchainImages;
    VkFormat swapchainFormat;
    VkExtent2D swapchainExtent;

    VkCommandPool commandPool;
    std::vector<VkCommandBuffer> commandBuffers;

    VkSemaphore imageAvailableSemaphore;
    VkSemaphore renderFinishedSemaphore;
};

VulkanContext vkContext;

// error checking macro
#define VK_CHECK_RESULT(f) \
    do { \
        VkResult res = (f); \
        if(res != VK_SUCCESS) { \
            __android_log_print(ANDROID_LOG_ERROR, "Vulkan", \
                               "Fatal error: VkResult = %d at %s:%d", \
                               res, __FILE__, __LINE__); \
            assert(res == VK_SUCCESS); \
        } \
    } while (0)

extern "C" JNIEXPORT void JNICALL
Java_com_june_myfps1_MainActivity_initVulkan(
        JNIEnv* env,
        jobject /* this */,
        jobject surface) {

    // 1. Instance Creation
    VkApplicationInfo appInfo{};
    appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
    appInfo.pApplicationName = "Vulkan App";
    appInfo.apiVersion = VK_API_VERSION_1_1;

    const char* extensions[] = {
            VK_KHR_SURFACE_EXTENSION_NAME,
            VK_KHR_ANDROID_SURFACE_EXTENSION_NAME
    };

    VkInstanceCreateInfo instanceCreateInfo{};
    instanceCreateInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    instanceCreateInfo.pApplicationInfo = &appInfo;
    instanceCreateInfo.enabledExtensionCount = 2;
    instanceCreateInfo.ppEnabledExtensionNames = extensions;

    VK_CHECK_RESULT(vkCreateInstance(&instanceCreateInfo, nullptr, &vkContext.instance));

    // 2. Surface Creation
    VkAndroidSurfaceCreateInfoKHR surfaceCreateInfo{};
    surfaceCreateInfo.sType = VK_STRUCTURE_TYPE_ANDROID_SURFACE_CREATE_INFO_KHR;
    surfaceCreateInfo.window = ANativeWindow_fromSurface(env, surface);
    VK_CHECK_RESULT(vkCreateAndroidSurfaceKHR(vkContext.instance, &surfaceCreateInfo,
                                              nullptr, &vkContext.surface));

    // 3. Physical Device Selection
    uint32_t deviceCount = 0;
    vkEnumeratePhysicalDevices(vkContext.instance, &deviceCount, nullptr);
    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(vkContext.instance, &deviceCount, devices.data());

    for (const auto& device : devices) {
        // Find queue families
        uint32_t queueFamilyCount = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);
        std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
        vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

        for (uint32_t i = 0; i < queueFamilyCount; i++) {
            VkBool32 presentSupport = false;
            vkGetPhysicalDeviceSurfaceSupportKHR(device, i, vkContext.surface, &presentSupport);

            if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT && presentSupport) {
                vkContext.graphicsQueueFamily = i;
                vkContext.presentQueueFamily = i;
                break;
            }
        }

        if (vkContext.graphicsQueueFamily != UINT32_MAX) {
            vkContext.physicalDevice = device;
            break;
        }
    }

    // 4. Logical Device Creation
    float queuePriority = 1.0f;
    VkDeviceQueueCreateInfo queueCreateInfo{};
    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queueCreateInfo.queueFamilyIndex = vkContext.graphicsQueueFamily;
    queueCreateInfo.queueCount = 1;
    queueCreateInfo.pQueuePriorities = &queuePriority;

    const char* deviceExtensions[] = { VK_KHR_SWAPCHAIN_EXTENSION_NAME };

    VkDeviceCreateInfo deviceCreateInfo{};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = 1;
    deviceCreateInfo.pQueueCreateInfos = &queueCreateInfo;
    deviceCreateInfo.enabledExtensionCount = 1;
    deviceCreateInfo.ppEnabledExtensionNames = deviceExtensions;

    VK_CHECK_RESULT(vkCreateDevice(vkContext.physicalDevice, &deviceCreateInfo,
                                   nullptr, &vkContext.device));

    vkGetDeviceQueue(vkContext.device, vkContext.graphicsQueueFamily, 0, &vkContext.graphicsQueue);
    vkGetDeviceQueue(vkContext.device, vkContext.presentQueueFamily, 0, &vkContext.presentQueue);

    // 5. Create Swapchain
    VkSurfaceCapabilitiesKHR surfaceCapabilities;
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(vkContext.physicalDevice, vkContext.surface,
                                              &surfaceCapabilities);

    uint32_t formatCount;
    vkGetPhysicalDeviceSurfaceFormatsKHR(vkContext.physicalDevice, vkContext.surface,
                                         &formatCount, nullptr);
    std::vector<VkSurfaceFormatKHR> surfaceFormats(formatCount);
    vkGetPhysicalDeviceSurfaceFormatsKHR(vkContext.physicalDevice, vkContext.surface,
                                         &formatCount, surfaceFormats.data());

    VkSurfaceFormatKHR chosenFormat = surfaceFormats[0];
    for (const auto& format : surfaceFormats) {
        if (format.format == VK_FORMAT_B8G8R8A8_SRGB) {
            chosenFormat = format;
            break;
        }
    }

    VkExtent2D extent = surfaceCapabilities.currentExtent;
    if (extent.width == UINT32_MAX) {
        // Handle window resize properly in real app
        extent.width = 1280;
        extent.height = 720;
    }

    VkSwapchainCreateInfoKHR swapchainCreateInfo{};
    swapchainCreateInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    swapchainCreateInfo.surface = vkContext.surface;
    swapchainCreateInfo.minImageCount = 2;
    swapchainCreateInfo.imageFormat = chosenFormat.format;
    swapchainCreateInfo.imageColorSpace = chosenFormat.colorSpace;
    swapchainCreateInfo.imageExtent = extent;
    swapchainCreateInfo.imageArrayLayers = 1;
    swapchainCreateInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
    swapchainCreateInfo.preTransform = surfaceCapabilities.currentTransform;
    swapchainCreateInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    swapchainCreateInfo.presentMode = VK_PRESENT_MODE_FIFO_KHR;

    VK_CHECK_RESULT(vkCreateSwapchainKHR(vkContext.device, &swapchainCreateInfo,
                                         nullptr, &vkContext.swapchain));

    vkGetSwapchainImagesKHR(vkContext.device, vkContext.swapchain,
                            &swapchainCreateInfo.minImageCount, nullptr);
    vkContext.swapchainImages.resize(swapchainCreateInfo.minImageCount);
    vkGetSwapchainImagesKHR(vkContext.device, vkContext.swapchain,
                            &swapchainCreateInfo.minImageCount, vkContext.swapchainImages.data());

    // 6. Command Pool & Buffers
    VkCommandPoolCreateInfo commandPoolCreateInfo{};
    commandPoolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    commandPoolCreateInfo.queueFamilyIndex = vkContext.graphicsQueueFamily;
    VK_CHECK_RESULT(vkCreateCommandPool(vkContext.device, &commandPoolCreateInfo,
                                        nullptr, &vkContext.commandPool));

    vkContext.commandBuffers.resize(swapchainCreateInfo.minImageCount);
    VkCommandBufferAllocateInfo commandBufferAllocateInfo{};
    commandBufferAllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    commandBufferAllocateInfo.commandPool = vkContext.commandPool;
    commandBufferAllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    commandBufferAllocateInfo.commandBufferCount = (uint32_t)vkContext.commandBuffers.size();
    VK_CHECK_RESULT(vkAllocateCommandBuffers(vkContext.device, &commandBufferAllocateInfo,
                                             vkContext.commandBuffers.data()));

    // 7. Create Semaphores
    VkSemaphoreCreateInfo semaphoreCreateInfo{};
    semaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
    VK_CHECK_RESULT(vkCreateSemaphore(vkContext.device, &semaphoreCreateInfo, nullptr,
                                      &vkContext.imageAvailableSemaphore));
    VK_CHECK_RESULT(vkCreateSemaphore(vkContext.device, &semaphoreCreateInfo, nullptr,
                                      &vkContext.renderFinishedSemaphore));
}

extern "C" JNIEXPORT void JNICALL
Java_com_june_myfps1_MainActivity_drawFrame(
        JNIEnv* /* env */,
        jobject /* this */) {

    uint32_t imageIndex;
    vkAcquireNextImageKHR(vkContext.device, vkContext.swapchain, UINT64_MAX,
                          vkContext.imageAvailableSemaphore, VK_NULL_HANDLE, &imageIndex);

    // Record command buffer
    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    VK_CHECK_RESULT(vkBeginCommandBuffer(vkContext.commandBuffers[imageIndex], &beginInfo));

    VkClearColorValue clearColor = {{0.2f, 0.9f, 0.8f, 1.0f}};
    VkImageSubresourceRange range{};
    range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    range.levelCount = 1;
    range.layerCount = 1;

    vkCmdClearColorImage(vkContext.commandBuffers[imageIndex], vkContext.swapchainImages[imageIndex],
                         VK_IMAGE_LAYOUT_PRESENT_SRC_KHR, &clearColor, 1, &range);

    VK_CHECK_RESULT(vkEndCommandBuffer(vkContext.commandBuffers[imageIndex]));

    // Submit to queue
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.waitSemaphoreCount = 1;
    submitInfo.pWaitSemaphores = &vkContext.imageAvailableSemaphore;
    VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
    submitInfo.pWaitDstStageMask = waitStages;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &vkContext.commandBuffers[imageIndex];
    submitInfo.signalSemaphoreCount = 1;
    submitInfo.pSignalSemaphores = &vkContext.renderFinishedSemaphore;

    VK_CHECK_RESULT(vkQueueSubmit(vkContext.graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE));

    // Present
    VkPresentInfoKHR presentInfo{};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    presentInfo.waitSemaphoreCount = 1;
    presentInfo.pWaitSemaphores = &vkContext.renderFinishedSemaphore;
    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = &vkContext.swapchain;
    presentInfo.pImageIndices = &imageIndex;

    VK_CHECK_RESULT(vkQueuePresentKHR(vkContext.presentQueue, &presentInfo));
}

extern "C" JNIEXPORT void JNICALL
Java_com_june_myfps1_MainActivity_cleanup(
        JNIEnv* /* env */,
        jobject /* this */) {

    vkDeviceWaitIdle(vkContext.device);

    vkDestroySemaphore(vkContext.device, vkContext.imageAvailableSemaphore, nullptr);
    vkDestroySemaphore(vkContext.device, vkContext.renderFinishedSemaphore, nullptr);
    vkDestroyCommandPool(vkContext.device, vkContext.commandPool, nullptr);
    vkDestroySwapchainKHR(vkContext.device, vkContext.swapchain, nullptr);
    vkDestroyDevice(vkContext.device, nullptr);
    vkDestroySurfaceKHR(vkContext.instance, vkContext.surface, nullptr);
    vkDestroyInstance(vkContext.instance, nullptr);
}